package zip;

import static zip.Utils.to;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.util.PriorityQueue;
import java.util.Queue;

public class HuffZip {
    private static Queue<Node> pq = null; // priorityQueue用于构造Huffman树时排序
    private static int[] freq = null; // 记录每个字符的出现次数
    private static String[] codes = null; // 记录Huffman编码
    private static Node root = null; // Huffman树根
    private static int bytesCount; // 用于记录一共有多少有效字符
    private static final int CHS = 300; // 总的字符种类有256个

    public static class Node implements Comparable<Node> {
        Node right, left;
        String code;
        int freq, ascii;

        public Node(Node l, Node r, int freq) {
            this.right = r;
            this.left = l;
            this.freq = freq;
        }

        @Override
        public int compareTo(Node o) {
            return this.freq - o.freq;
        }
    }

    /**
     * 压缩文件对外的接口方法
     *
     * @param pathName 要压缩的文件路径
     */
    public static void doHuffZip(String pathName) {
        initZip();
        calFreq(pathName);
        createTree();
        getCode(root, "");
        realZip(midZip(pathName), pathName + ".huff");
    }

    private static void initZip() {
        pq = new PriorityQueue<>();
        freq = new int[CHS];
        codes = new String[CHS];
        root = null;
        bytesCount = 0;
    }

    /**
     * 统计文本每一个byte的出现次数
     *
     * @param pathName 输入文本路径名
     */
    private static void calFreq(String pathName) {
        File file = new File(pathName);
        try {
            DataInputStream din = new DataInputStream(new FileInputStream(file));
            while (true) {
                try {
                    freq[to(din.readByte())] += 1;
                } catch (EOFException eof) {
                    break;
                }
            }
            din.close();
        } catch (IOException ioe) {
            System.out.println("IO Exception = " + ioe);
        }
    }

    /**
     * 构建huffman树
     */
    private static void createTree() {
        for (int i = 0; i < CHS; ++i) {
            if (freq[i] > 0) {
                Node node = new Node(null, null, freq[i]);
                node.ascii = i;
                pq.add(node);
                bytesCount += 1;
            }
        }
        if (bytesCount == 0)
            return;
        while (pq.size() > 1) {
            Node temp1 = pq.poll(), temp2 = pq.poll();
            Node fath = new Node(temp1, temp2, temp1.freq + temp2.freq);
            pq.add(fath);
        }
        root = pq.poll();
    }

    /**
     * 前序遍历树，获取字符的哈夫曼编码
     *
     * @param n 树根
     * @param code 当前构建的哈夫曼编码
     */
    private static void getCode(Node n, String code) {
        if (n == null)
            return;
        n.code = code;
        if (n.left == null && n.right == null)
            codes[n.ascii] = code;
        getCode(n.left, code + "0");
        getCode(n.right, code + "1");
    }

    /**
     * 压缩时临时产生辅助的中间文件 </br>
     * 文件中以字符串的形式记录哈夫曼编码 所以物理上的一个byte存储逻辑上的一个bit
     *
     * @param pathName 要压缩的文件路径
     * @return 中间文件的路径名
     */
    private static String midZip(String pathName) {
        String midpath = pathName + ".mid";
        File midFile = new File(midpath);
        try {
            DataInputStream din = new DataInputStream(new FileInputStream(new File(pathName)));
            PrintStream printStream = new PrintStream(midFile);

            while (true) {
                try {
                    byte bt = din.readByte();
                    printStream.print(codes[to(bt)]);
                } catch (EOFException eof) {
                    break;
                }
            }
            din.close();
            printStream.close();
        } catch (IOException ioe) {
            System.out.println("IO Exception = " + ioe);
        }
        return midpath;
    }

    /**
     * 将01字符串压缩成01位串
     *
     * @param midPath 中间文件，存储的是01字符串
     * @param tarPath 目标输出文件，存储压缩之后的结果
     * @apiNote 目标文件的存储格式：</br>
     * 总的有效字符数[int] + (有效字符[byte] + 频数[int])s +</br>
     * 多余的凑数bit个数[int] + 哈夫曼编码[bits]s
     */
    private static void realZip(String midPath, String tarPath) {
        File inFile = new File(midPath);
        File outFile = new File(tarPath);
        try {
            DataInputStream din = new DataInputStream(new FileInputStream(inFile));
            DataOutputStream dout = new DataOutputStream(new FileOutputStream(outFile));

            dout.write(bytesCount);
            // System.out.println("bytes counter = " + bytesCount);
            for (int i = 0; i < (1 << 8); ++i) {
                if (freq[i] > 0) {
                    byte ch = (byte) i;
                    dout.write(ch);
                    dout.writeInt(freq[i]);
                }
            }

            byte extraBitsCount = (byte) (int) ((8L - (inFile.length() % 8L)) % 8L);
            dout.write(extraBitsCount);
            // System.out.println("extra bits count = " + extraBitsCount);

            // 将01字符串压缩成位串
            // 由于Java中最小能够直接表示的类型为byte
            // 所以利用位操作
            byte bits = 0, ch;
            int iter;
            while (true) {
                try {
                    for (iter = 0; iter < 8; ++iter) {
                        ch = din.readByte();
                        bits <<= 1;
                        if (ch == '1')
                            bits |= 1;
                    }
                    dout.write(bits);
                } catch (EOFException eof) {
                    if (extraBitsCount > 0) {
                        bits <<= extraBitsCount;
                        dout.write(bits);
                    }
                    break;
                }
            }
            din.close();
            dout.close();
        } catch (IOException ioe) {
            System.out.println("IO Exception = " + ioe);
        }
        inFile.delete(); // 删除中间文件
    }
}
